# -*- coding: utf-8 -*-

from pylab import *
from math import fsum
from scipy import signal

# scaling function coefficients
#h = array([(1 + 3**0.5) / (4 * 2**0.5),
#     (3 + 3**0.5) / (4 * 2**0.5),
#     (3 - 3**0.5) / (4 * 2**0.5),
#     (1 - 3**0.5) / (4 * 2**0.5)])
# wavelet function coefficients
h = [0.6830127, 1.1830127, 	0.3169873, -0.1830127]
g = array([h[3], -h[2], h[1], -h[0]])
hlast = [h[0] - h[2] - 2 * h[3], h[1] + 2 * h[2] + 3 * h[3]]
glast = [g[0] - g[2] - 2 * g[3], g[1] + 2 * g[2] + 3 * g[3]]

def mxrange(start, stop, mod):
    for i in xrange(start, stop):
        yield i % mod

def printout(i, x):
    print i, ' '.join([str(n) for n in x])

def fwtd(x):
    n = int(log2(x.size))
    tx = array(x) * 1.0
    
    for i in xrange(1, n):
        ttx = array(tx)
        jMax = 1 << (n - i)
        for j in xrange(0, jMax):
            l = list(mxrange(2 * j, 2 * j + 4, 2 * jMax))
            r0 = j
            r1 = r0 + jMax
            ttx[r0] = fsum(tx[i] * h[k] for i, k in zip(l, xrange(4))) / 2
            ttx[r1] = fsum(tx[i] * g[k] for i, k in zip(l, xrange(4))) / 2
        tx = array(ttx)
    return tx

def ifwtd(x):
    n = int(log2(x.size))
    tx = array(x) * 1.0
    
    for i in xrange(n - 1, 0, -1):
        ttx = zeros(x.size)
        jMax = 1 << (n - i)
        for j in xrange(jMax - 1, -1, -1):
            l0, l1 = j, j + jMax
            r = mxrange(2 * j, 2 * j + 4, 2 * jMax)
            for k, m in zip(r, xrange(4)):
                ttx[k] += tx[l0] * h[m] + tx[l1] * g[m]
        # записываем неизменённые компоненты
        for j in xrange(2 * jMax, x.size):
            ttx[j] = tx[j]
        tx = array(ttx)
    return tx    

if __name__ == '__main__':
    #t = linspace(0, 7, 8)
    #sig = signal.sawtooth(t, 0.8)
    #t = linspace(0, 8 * 512 - 1, 8 * 512)
    #sig = repeat(sig, 512)
    #sig += randn(len(sig))
    t = linspace(0, 127, 128)
    sig = array(t)
    
    X = fwtd(sig)
    figure()
    subplot(3, 1, 1)
    plot(t, sig)
    subplot(3, 1, 2)
    plot(t, X)
    X[1024:len(X)] = 0
    subplot(3, 1, 3)
    plot(t, sig, t, ifwtd(X), 'r')
    show()
